import { useEffect, useRef } from 'react';

export const useDeviceWebSocket = (deviceId: string, onDataReceived?: (data: string) => void) => {
  const socketRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    if (!deviceId) return;
    
    const socket = new WebSocket(`ws://localhost:8000/devices/screen_stream_v2/${deviceId}`);
    socketRef.current = socket;

    socket.onopen = () => {
      console.log('WebSocket连接已建立');
    };

    socket.onmessage = (event) => {
      if (typeof onDataReceived === 'function') {
        // 将二进制数据转换为base64字符串
        const reader = new FileReader();
        reader.onload = () => {
          const base64data = reader.result as string;
          onDataReceived(base64data.split(',')[1]); // 提取base64部分
        };
        reader.readAsDataURL(new Blob([event.data], {type: 'image/jpeg'}));
      }
    };

    socket.onclose = () => {
      console.log('WebSocket连接已关闭');
    };

    return () => {
      if (socketRef.current) {
        socketRef.current.close();
      }
    };
  }, [deviceId]);

  return socketRef.current;
};